import torch as pt
import torchvision as ptv


class Resnet34TransModel(pt.nn.Module):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)

        resnet = ptv.models.resnet34(pretrained=True)
        resnet_out = resnet.fc.in_features
        self.base_model = pt.nn.Sequential(*(list(resnet.children())[:-1]))
        for p in self.base_model.parameters():
            p.requires_grad = False

        self.customer_model = pt.nn.Linear(resnet_out, n_cls)

    def forward(self, x):
        x = self.base_model(x)
        x = pt.squeeze(x, dim=3)
        x = pt.squeeze(x, dim=2)
        x = self.customer_model(x)
        return x


model = Resnet34TransModel(2)
print(model)
x = pt.ones(4, 3, 224, 224, dtype=pt.float32)
pred = model(x)
print(pred)
